import matplotlib.pyplot as plt
import networkx as nx
from collections import Counter
import itertools
import numpy as np
import random
from collections import deque
import math
import os


def generate_test_graph(
    n_nodes=100,
    k_dist=4,
    target_edge=None,
    seed=None,
):
    """
    生成一个周期性(toroidal)的二维正方格子图, 并标记所有与 target_edge 距离正好 k_dist 的边.

    Parameters
    ----------
    n_nodes : int
        节点总数, 必须是完全平方数 (如 400 = 20×20)
    target_edge : tuple or None
        如果给定, 用作“中心”的那条边; 否则随机选一条.
        格式为 (u, v), 其中 u, v 是节点整数标签.
    k_dist : int
        要标记的距离阈值
    seed : int or None
        随机种子, 便于复现

    Returns
    -------
    G : networkx.Graph
        带有以下属性的图:
          - 每条边有布尔属性 'break_flag'
          - 每个节点有位置属性 'position'
    """
    if seed is not None:
        random.seed(seed)

    # 1) 生成 m×m 的周期性二维正方格子
    m = int(math.isqrt(n_nodes))
    if m * m != n_nodes:
        raise ValueError(f"n_nodes={n_nodes} 不是完全平方数, 无法生成正方格子")
    # 生成带周期边的 2D 网格,节点标签是 (i,j)
    G_grid = nx.grid_2d_graph(m, m, periodic=False)

    # 把 (i,j) 标签换成整数 0..n_nodes-1,并保留位置
    mapping = {(i, j): i * m + j for i, j in G_grid.nodes()}
    G = nx.relabel_nodes(G_grid, mapping)
    # 保存位置属性,供可视化时使用
    pos = {mapping[(i, j)]: (j, -i) for i, j in G_grid.nodes()}
    nx.set_node_attributes(G, pos, 'position')

    # 2) 确定 target_edge
    edges = list(G.edges())
    if target_edge is None:
        target_edge = random.choice(edges)
    elif target_edge not in edges and (target_edge[1], target_edge[0]) not in edges:
        raise ValueError("你指定的 target_edge 不在 G 中")
    # 统一为 (u, v) 且 u<v
    target_edge = tuple(sorted(target_edge))

    # 3) 构造 Line Graph 并 BFS 计算每条边的距离
    L = nx.line_graph(G)
    dist = {edge: None for edge in edges}
    dq = deque([target_edge])
    dist[target_edge] = 0

    while dq:
        curr = dq.popleft()
        d = dist[curr]
        if d >= k_dist:
            continue
        for nei in L.neighbors(curr):
            if dist[nei] is None:
                dist[nei] = d + 1
                dq.append(nei)

    # 4) 标记 distance == k_dist 的边
    nx.set_edge_attributes(G, False, 'break_flag')
    break_count = 0
    for e, d in dist.items():
        if d == k_dist:
            G.edges[e]['break_flag'] = True
            break_count += 1

    print(f"Number of edges with break_flag=True: {break_count}")
    return G

# 只标记两个距离确定的 edge


def generate_test_graph_1(
    n_nodes: int = 100,
    k_dist: int = 4,
    seed: int = None,
) -> nx.Graph:
    """
    生成一个 m×m 周期性(torus)正方格子图, 并在两条恰好 X 方向相隔 k_dist 的“平行”边上
    标记 'break_flag'=True.所有节点以 (i,j) 坐标对表示,均在第一象限 (>=0).

    具体地,选定中心边 ((0,0),(0,1)),再向右水平平移 k_dist 后取对应平行边:
    ((0,k_dist),(0,(k_dist+1) mod m)).

    Parameters
    ----------
    n_nodes : int
        节点总数, 必须是完全平方数
    k_dist : int
        水平方向(X 方向)的格子距离 (1 ≤ k_dist ≤ m-1)
    seed : int or None
        随机种子 (不影响本函数的边选择,因为我们固定了 target_edge)

    Returns
    -------
    G : networkx.Graph
        - 节点以 (i,j) 二元组表示
        - 每条边带布尔属性 'break_flag'
        - 每个节点带位置属性 'position',且坐标均 ≥ 0
    """
    # 1) 随机种子
    if seed is not None:
        random.seed(seed)

    # 2) 构造 m×m 周期性网格 (torus),节点即为 (i,j)
    m = int(math.isqrt(n_nodes))
    if m * m != n_nodes:
        raise ValueError(f"n_nodes={n_nodes} 不是完全平方数")

    G = nx.grid_2d_graph(m, m, periodic=True)

    # 3) 检查 k_dist 合法性
    if not (1 <= k_dist <= m - 1):
        raise ValueError(f"k_dist={k_dist} 超出允许范围 [1, {m-1}]")

    # 4) 初始化所有边的 break_flag=False
    nx.set_edge_attributes(G, False, 'break_flag')

    # 5) 标记第一条边 ((0,0),(0,1))
    u = (0, 0)
    v = (0, 1)
    G.edges[u, v]['break_flag'] = True

    # 6) 计算并标记第二条边 ((0,k_dist),(0,(k_dist+1) mod m))
    j_start = k_dist % m
    j_end = (k_dist + 1) % m
    u2 = (0, j_start)
    v2 = (0, j_end)
    G.edges[u2, v2]['break_flag'] = True

    # 7) 为每个节点设置位置属性
    pos = {node: node for node in G.nodes()}
    nx.set_node_attributes(G, pos, 'position')

    print(f"Marked edges: {u}-{v} and {u2}-{v2}")
    return G

# 随机抽取一定比例的 edge 标记


def generate_test_graph_2(n_nodes=100, proportion=0.1,
                          periodic=True, attr_name='break_flag'):
    """
    从图 G 中随机抽取一定比例的边进行标记.

    参数
    ----
    G : networkx.Graph
        待操作的图.
    proportion : float
        需要抽取的边的比例,取值范围 [0, 1].
    attr_name : str, optional
        边属性的名称,默认 'marked'.

    返回
    ----
    sampled_edges : list of tuple
        被标记的边列表,每个元素是 (u, v).
    """
    # 1) 构造 m×m 网格
    m = int(math.isqrt(n_nodes))
    if m*m != n_nodes:
        raise ValueError(f"n_nodes={n_nodes} 不是完全平方数")
    G0 = nx.grid_2d_graph(m, m, periodic=periodic)
    # 重标号为 0..n_nodes-1,并设置可视化用位置
    mapping = {(i, j): i*m + j for i, j in G0.nodes()}
    G = nx.relabel_nodes(G0, mapping)
    pos = {mapping[(i, j)]: (j, -i) for i, j in G0.nodes()}
    nx.set_node_attributes(G, pos, 'position')

    if not (0 <= proportion <= 1):
        raise ValueError(f"proportion 必须在 [0,1] 之间,当前为 {proportion}")

    edges = list(G.edges())
    n_edges = len(edges)
    n_select = int(round(proportion * n_edges))

    # 随机抽样
    sampled_edges = random.sample(edges, n_select)

    # 先将所有边的属性初始化为 False
    nx.set_edge_attributes(G, False, attr_name)
    # 对抽中的边打标记
    for u, v in sampled_edges:
        G[u][v][attr_name] = True

    return G


def visualize_graph_old(
    G,
    node_size=20,
    edge_width=1.0,
    break_edge_width=2.0,
    break_color='red',
    default_color='lightgray',
    figsize=(8, 8),
    save_path=None
):
    """
    可视化图 G:节点显示为小圆点,所有 break_flag=True 的边高亮.

    Parameters
    ----------
    G : networkx.Graph
        节点必须有 'position' 属性,边必须有 'break_flag' 属性.
    node_size : float
        节点大小
    edge_width : float
        普通边的宽度
    break_edge_width : float
        break_flag边的宽度
    break_color : color
        break_flag边的颜色
    default_color : color
        其他边的颜色
    figsize : tuple
        图像大小
    save_path : str or None
        如果指定,则保存到该路径,否则直接 plt.show()
    """
    pos = nx.get_node_attributes(G, 'position')
    # 分离两类边
    default_edges = [(u, v) for u, v, d in G.edges(
        data=True) if not d.get('break_flag', False)]
    break_edges = [(u, v) for u, v, d in G.edges(
        data=True) if d.get('break_flag', False)]

    plt.figure(figsize=figsize)
    # 先画普通边
    nx.draw_networkx_edges(
        G, pos,
        edgelist=default_edges,
        width=edge_width,
        edge_color=default_color
    )
    # 再叠加高亮边
    nx.draw_networkx_edges(
        G, pos,
        edgelist=break_edges,
        width=break_edge_width,
        edge_color=break_color
    )
    # 最后画节点
    nx.draw_networkx_nodes(G, pos, node_size=node_size, node_color='black')

    plt.axis('off')
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()


def visualize_graph(
    G,
    node_size=20,
    edge_width=1.0,
    break_edge_width=2.0,
    break_color='red',
    default_color='lightgray',
    figsize=(8, 8),
    box_size=None,     # 新增:周期盒子大小 (Lx, Ly),不传就按普通方式画
    save_path=None
):
    pos = nx.get_node_attributes(G, 'position')
    default_edges = [(u, v) for u, v, d in G.edges(
        data=True) if not d.get('break_flag', False)]
    break_edges = [(u, v) for u, v, d in G.edges(
        data=True) if d.get('break_flag', False)]
    plt.figure(figsize=figsize)
    if box_size is None:
        # 和你原来的一样
        nx.draw_networkx_edges(
            G, pos, edgelist=default_edges, width=edge_width,   edge_color=default_color)
        nx.draw_networkx_edges(G, pos, edgelist=break_edges,
                               width=break_edge_width, edge_color=break_color)
    else:
        Lx, Ly = box_size
        import numpy as _np

        import numpy as _np

        def edge_segments(u, v, pos, Lx, Ly):
            """
            给定节点 u, v;pos 是 dict:node -> (x,y);Lx, Ly 是盒子长度
            返回一条边在 PBC 下应该画的 1~2 段线段,每段 [(x1,y1),(x2,y2)].
            """
            x1, y1 = pos[u]
            x2, y2 = pos[v]
            dx = x2 - x1
            dy = y2 - y1

            # 最小镜像:把 dx, dy 压到 [-L/2, L/2)
            if dx > Lx/2:
                dx -= Lx
            elif dx < -Lx/2:
                dx += Lx
            if dy > Ly/2:
                dy -= Ly
            elif dy < -Ly/2:
                dy += Ly
            x2s = x1 + dx
            y2s = y1 + dy

            # 收集所有"可能的打断点"(t, axis, boundary_value)
            cuts = []
            # x 方向边界
            if dx != 0:
                for xb in (0.0, Lx):
                    t = (xb - x1) / dx
                    # 只要在 (0,1) 之间,才算真的有穿
                    if 0 < t < 1:
                        cuts.append((t, 'x', xb))
            # y 方向边界
            if dy != 0:
                for yb in (0.0, Ly):
                    t = (yb - y1) / dy
                    if 0 < t < 1:
                        cuts.append((t, 'y', yb))

            # 没有任何穿边界点,就画一整段
            if not cuts:
                return [[(x1, y1), (x2s, y2s)]]

            # 找最先穿的那个
            t0, axis, bound = min(cuts, key=lambda x: x[0])
            # 计算真正的交点
            xc = x1 + dx * t0
            yc = y1 + dy * t0

            segs = []
            # 第一段:从 (x1,y1) 到交点 (xc,yc)
            segs.append([(x1, y1), (xc, yc)])

            # 第二段:从交点“跳”到对侧再画到 (x2s,y2s)
            if axis == 'x':
                # x 轴周期跳转
                x_jump = xc - _np.sign(dx) * Lx
                segs.append([(x_jump, yc), (x2s, y2s)])
            else:
                # y 轴周期跳转
                y_jump = yc - _np.sign(dy) * Ly
                segs.append([(xc, y_jump), (x2s, y2s)])

            return segs

        # 用 matplotlib 逐条画
        for u, v in default_edges:
            for (p1, p2) in edge_segments(u, v, pos, Lx, Ly):
                plt.plot([p1[0], p2[0]], [p1[1], p2[1]],
                         lw=edge_width, color=default_color, zorder=1, solid_capstyle='round')
        for u, v in break_edges:
            for (p1, p2) in edge_segments(u, v, pos, Lx, Ly):
                plt.plot([p1[0], p2[0]], [p1[1], p2[1]],
                         lw=break_edge_width, color=break_color,   zorder=2, solid_capstyle='round')
    # 最后画节点
    node_collection = nx.draw_networkx_nodes(
        G, pos,
        node_size=node_size,
        node_color='black'
    )
    node_collection.set_zorder(3)
    plt.axis('off')
    if not save_path:
        save_path = './graph_visualize.png'
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()


def plot_grid_network(G, edge_gebc, Nx, Ny, step, norm=None):
    """绘制二维格点网络,并根据 GEBC 值对边进行着色,排除周期性边界的边"""

    output_folder = "frames"

    # 获取实际存在于图 G 中的边,并保证边的节点顺序一致
    edges = list(G.edges())
    sorted_edges = [tuple(sorted(edge)) for edge in edges]
    gebc_values = [edge_gebc[edge] for edge in sorted_edges]
    rounded_gebc_values = [value for value in gebc_values]

    # # 规范化 GEBC 值以便颜色映射
    # gebc_colors = np.array(rounded_gebc_values)
    # norm = plt.Normalize(vmin=min(gebc_colors), vmax=max(gebc_colors))
    # cmap = plt.cm.viridis

    # 将 GEBC 值转换为数组
    gebc_colors = np.array(rounded_gebc_values)
    max_abs = max(abs(gebc_colors.min()), abs(gebc_colors.max()))
    # 在外面已经处理好了
    norm = plt.Normalize(vmin=-max_abs, vmax=max_abs)
    cmap = plt.cm.coolwarm    # 选择双向颜色映射

    # 绘图
    pos = {(x, y): (x, y) for x, y in G.nodes()}  # 节点的位置为其坐标
    plt.figure(figsize=(10, 10))

    # 绘制非周期性边界的边
    for i, edge in enumerate(edges):
        node1, node2 = edge

        # 判断是否为周期性边界的边
        if (abs(node1[0] - node2[0]) == Nx - 1 and node1[1] == node2[1]) or \
           (abs(node1[1] - node2[1]) == Ny - 1 and node1[0] == node2[0]):
            continue  # 跳过周期性边界的边

        # 绘制边
        x_values = [pos[node1][0], pos[node2][0]]
        y_values = [pos[node1][1], pos[node2][1]]
        plt.plot(
            x_values, y_values,
            color=cmap(norm(gebc_colors[i])),
            linewidth=4,
            zorder=1  # 设置边的绘制顺序低于节点
        )

    # # 绘制节点
    # node_x = [pos[node][0] for node in G.nodes()]
    # node_y = [pos[node][1] for node in G.nodes()]
    # plt.scatter(node_x, node_y, s=30, color='black', zorder=0)  # 增大节点尺寸,使其完全覆盖边界
    # 定义颜色条的映射器
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])

    # 在这里使用 `plt.gca()`,获取当前的轴,并将其传入 `plt.colorbar`
    cbar = plt.colorbar(sm, ax=plt.gca(), label="Edge GEBC (Non-normalized)")

    # plt.title(f"Step {step + 1} - Remove Edge {edge_to_remove}")
    plt.title(f"Step {step + 1}")
    plt.axis('equal')

    # 创建输出文件夹(如果不存在)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 保存帧图像
    plt.savefig(f'{output_folder}/frame_{step + 1}.png')
    plt.close()
